import json
import string
from operator import ge
import cv2
from cv2 import matchShapes
import numpy as np
from PIL import Image
import random
def get_point(key_points,index):
    return (int(key_points[index*3]),int(key_points[index*3+1]))
def rm_bg(id):
    img = cv2.imread('./cloth/0'+id+'_cloth.jpg', cv2.IMREAD_UNCHANGED)
    # img = cv2.resize(img,(450,800))
    mask = np.zeros(img.shape[:2], dtype=np.uint8)

    # 矩形roi
    rect = (43,125,700,750)  # 包括前景的矩形，格式为(x,y,w,h)

    bgdmodel = np.zeros((1, 65), np.float64)  # bg模型的临时数组
    fgdmodel = np.zeros((1, 65), np.float64)  # fg模型的临时数组

    cv2.grabCut(img, mask, rect, bgdmodel, fgdmodel, 1, mode=cv2.GC_INIT_WITH_RECT)

    # 提取前景和可能的前景区域
    mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
    img = img * mask2[:, :, np.newaxis]

    img = cv2.bitwise_and(img, img, mask=mask2)
    cv2.imwrite('./cloth/0'+id+'rmbg.jpg', img)

def change_cloth(id:string):
    rm_bg(id)
    # 首先读入img
    img = cv2.imread('./cloth/0'+id+'rmbg.jpg', cv2.IMREAD_UNCHANGED)
    # img = cv2.resize(img,(450,800))

    # 读取图片关键点信息
    with open('./graph/0'+id+'_keypoints.json') as f:
        line = f.readline()
        json_obj = json.loads(line)
        keypoints = json_obj["people"][0]["pose_keypoints_2d"]

    new_points = []
    new_points.append(get_point(keypoints, 1))
    new_points.append(get_point(keypoints, 2))
    # new_points.append(get_point(keypoints, 3))
    new_points.append(get_point(keypoints, 5))
    # new_points.append(get_point(keypoints, 6))
    new_points.append(get_point(keypoints, 8))
    new_points.append(get_point(keypoints, 9))
    new_points.append(get_point(keypoints, 12))

    # 读取衣物关键点信息

    with open("./cloth/0"+id+".json") as f:
        line = f.readline()
        json_obj = json.loads(line)
        keypoints = json_obj["keypoints"]

    points = []
    points.append(get_point(keypoints, 0))
    points.append(get_point(keypoints, 1))
    points.append(get_point(keypoints, 2))
    points.append(get_point(keypoints, 3))
    points.append(get_point(keypoints, 4))
    points.append(get_point(keypoints, 5))
    # points.append(get_point(keypoints, 6))
    # points.append(get_point(keypoints, 7))

    # 缩放，使源点与目标点大小统一
    sx = float(new_points[5][1] - new_points[0][1]) / (points[5][1] - points[0][1])
    img = cv2.resize(img, None, fx=sx, fy=sx, interpolation=cv2.INTER_AREA)
    for i in range(0, len(points)):
        points[i] = (int(points[i][0] * sx), int(points[i][1] * sx))

    # 周围拓宽一圈
    offset = 200
    img = cv2.copyMakeBorder(img, offset, offset, offset, offset, cv2.BORDER_REPLICATE)
    for i in range(0, len(points)):
        points[i] = (int(points[i][0] + offset), int(points[i][1] + offset))
        new_points[i] = (int(new_points[i][0] + offset), int(new_points[i][1] + offset))

    tx = points[0][0] - new_points[0][0]
    ty = points[0][1] - new_points[0][1]

    # 平移使点之间距离缩小
    w = img.shape[1]
    h = img.shape[0]
    moving_matrix = np.float64([[1, 0, -tx], [0, 1, -ty]])
    img = cv2.warpAffine(img, moving_matrix, (w, h))

    for i in range(0, len(points)):
        points[i] = (int(points[i][0] - tx), int(points[i][1] - ty))

    # 画出来看一下当前衣服的标记点位置
    # for point in points:
    # 	cv2.circle(img, point, 1, (0, 255, 0), 2)
    # cv2.imshow("name",img)

    # tps
    tps = cv2.createThinPlateSplineShapeTransformer()

    # 源点的矩阵
    sourceshape = np.array(points, np.int32)
    sourceshape = sourceshape.reshape(1, -1, 2)
    print(sourceshape)

    # 创建点之间的匹配关系
    matches = []
    for i in range(1, len(points) + 1):
        matches.append(cv2.DMatch(i, i, 0))

    print(points, new_points)

    # 目标点的矩阵
    targetshape = np.array(new_points, np.int32)
    targetshape = targetshape.reshape(1, -1, 2)
    print(targetshape)

    tps.estimateTransformation(targetshape, sourceshape, matches)
    img = tps.warpImage(img)
    cv2.imwrite("./cloth/tps/0"+id+".jpg",img)
    # for point in new_points:
    #     cv2.circle(img, point, 1, (0, 255, 0), 2)
    moving_matrix=np.float64([[1,0,-offset],[0,1,-offset]])
    img=cv2.warpAffine(img,moving_matrix,(w,h))
    bg=cv2.imread('add_background/0'+id+'.jpg')
    # cv2.imshow('l',img)
    img=img[0:bg.shape[0],0:bg.shape[1]]
    img=cv2.copyMakeBorder(img,0,bg.shape[0]-img.shape[0],0,bg.shape[1]-img.shape[1],cv2.BORDER_CONSTANT,value=[0,0,0])
    print(bg.shape[0]-img.shape[0],bg.shape[1]-img.shape[1])
    # 原图mask
    lower=np.array([0,0,0])
    upper=np.array([1,1,1])
    mask=cv2.inRange(img,lower,upper)
    masked_image = np.copy(img)
    crop_background = np.copy(bg)
    crop_background[mask == 0] = [0, 0, 0]
    complete_image = masked_image + crop_background

    # cv2.imshow('bg',bg)
    # cv2.imshow('res',complete_image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    cv2.imwrite('./res/0'+id+'_res.jpg',complete_image)

for i in range(1,6):
    change_cloth(str(i))
